package com.demo.component.ollama.generate.impl;

import com.demo.component.ollama.generate.inter.GenerateASync;
import io.github.ollama4j.OllamaAPI;
import io.github.ollama4j.exceptions.OllamaBaseException;
import io.github.ollama4j.models.response.OllamaAsyncResultStreamer;
import io.github.ollama4j.models.response.OllamaResult;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.io.IOException;

@Service("GenerateASyncImpl")
public class GenerateASyncImpl implements GenerateASync {

    @Value("${application.ollama.url}")
    private String url;

    @Value("${application.ollama.model}")
    private String model;

    @Override
    public String generateAnswer(String prompt) {

        OllamaAPI ollamaAPI = new OllamaAPI(url);

        OllamaAsyncResultStreamer streamer = ollamaAPI.generateAsync(model, prompt, false);

        int pollIntervalMilliseconds = 1000;

        StringBuffer answer = new StringBuffer();

        while (true) {
            String tokens = streamer.getStream().poll();
            answer.append(tokens);
            if (!streamer.isAlive()) {
                break;
            }
            try {
                Thread.sleep(pollIntervalMilliseconds);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }

        return answer.toString();
    }
}
